# 作者：宋安康
# 开发时间：2023/11/19 16:18
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import numpy as np

mnist = fetch_openml('mnist_784', version=1)
X = mnist['data']
y = mnist['target']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
dt = DecisionTreeClassifier(max_depth=None, random_state=42)
dt.fit(X_train, y_train)

y_pred = dt.predict(X_test)
score = accuracy_score(y_test, y_pred)
print("Accuracy:", score)

# 使用 matplotlib 和 sklearn.tree 的 plot_tree 方法来可视化决策树
fig, axes = plt.subplots(4, 4, figsize=(10, 10))
for i, ax in enumerate(axes.flat):
    plot_tree(dt, X_test[i].reshape(28, 28), ax=ax)  # 使用plot_tree方法绘制决策树
    ax.set_title(f"Prediction: {y_pred[i]}")  # 设置标题为预测结果而不是决策规则
plt.show()